import cv2
import numpy as np
import copy

from utils import img_util 

algo_id = 20

def draw_box_by_data(boxes, scores, label_id, label_dict, img, device):
    boxes, scores, label_id = np.squeeze(boxes), np.squeeze(scores), np.squeeze(label_id).astype(np.int32)
    save = False
    y_len, x_len = img.shape[:2]
    # box_info = {'ref_clothes_box': []}

    for (box, score, idx) in zip(boxes, scores, label_id):
        label = label_dict[idx]
        if score > 0.7 and label == 'no_ref':
            y_min, x_min, y_max, x_max = box
            pt1 = (int(x_min * x_len), int(y_min * y_len))
            pt2 = (int(x_max * x_len), int(y_max * y_len))
            cv2.rectangle(img, pt1, pt2, (0, 0, 255), thickness=5)
            # box_info['ref_clothes_box'].append((pt1, pt2, (0, 0, 255)))
            save = True

    # 对识别结果做去重操作
    if save:
        img_util.save_image(img, device, algo_id)

    return img


def draw_box_by_data_trt(boxes, scores, label_id, label_dict, img, device):
    #boxes, scores, label_id = np.squeeze(boxes.numpy()), np.squeeze(scores.numpy()), np.squeeze(label_id.numpy()).astype(np.int32)
    save = False
    #y_len, x_len = img.shape[:2]
    # box_info = {'ref_clothes_box': []}

    for (box, score, idx) in zip(boxes.tolist(), scores.tolist(), label_id.tolist()):
        c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
        label = label_dict[idx]
        if score > 0.65 and label == 'other_clothes':
            cv2.rectangle(img, c1, c2, (0, 0, 255), thickness=4)
            # box_info['ref_clothes_box'].append((pt1, pt2, (0, 0, 255)))
            save = True
    cv2.waitKey(1)
    # 对识别结果做去重操作
    if save:
        ori_img = copy.deepcopy(img)
        img_util.save_image(img, device, algo_id, ori_img)

    return img

